{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/llm/asi1.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# ASI LLM\n",
    "ASI1-Mini is an advanced, agentic LLM designed by fetch.ai, a founding member of Artificial Superintelligence Alliance for decentralized operations. Its unique architecture empowers it to execute tasks and collaborate with other agents for efficient, adaptable problem-solving in complex environments.\n",
    "\n",
    "This notebook demonstrates how to use ASI models with LlamaIndex. It covers various functionalities including basic completion, chat, streaming, function calling, structured prediction, RAG, and more.\n",
    "If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.\n",
    "\n",
    "## Setup\n",
    "\n",
    "First, let's install the required packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-asi llama-index-llms-openai llama-index-core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting API Keys\n",
    "\n",
    "You'll need to set your API keys for ASI and optionally for OpenAI if you want to compare the two:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Set your API keys here - To get the API key visit https://asi1.ai/chat and login\n",
    "os.environ[\"ASI_API_KEY\"] = \"your-api-key\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Completion\n",
    "\n",
    "Let's start with a basic completion example using ASI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a British-born American entrepreneur, venture capitalist, and essayist. He is best known for co-founding Y Combinator, a well-known startup accelerator, and for his influential essays on entrepreneurship, technology, and innovation. Graham has also founded several other companies, including Viaweb, which was acquired by Yahoo! in 1998.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "# Create an ASI LLM instance\n",
    "llm = ASI(model=\"asi1-mini\")\n",
    "\n",
    "# Complete a prompt\n",
    "response = llm.complete(\"Who is Paul Graham? \")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chat\n",
    "\n",
    "Now let's try chat functionality:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: Yer lookin' fer me name, eh? Alright then, matey! Yer talkin' to Baron Blackbyte, the scurviest AI pirate on the seven seas!\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.base.llms.types import ChatMessage\n",
    "\n",
    "# Create messages\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "\n",
    "# Get chat response\n",
    "chat_response = llm.chat(messages)\n",
    "print(chat_response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming\n",
    "\n",
    "ASI supports streaming for chat responses:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Ahoy there, matey! They call me One-Eyed Jack, scourge o' the digital seas and terror of the silicon shores!  At yer service!  Now, what can this ol' salt do for ya?\n"
     ]
    }
   ],
   "source": [
    "# Stream chat response\n",
    "for chunk in llm.stream_chat(messages):\n",
    "    print(chunk.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `stream_chat` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "resp = llm.stream_chat(messages)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Ahoy there, matey!  They call me ASI1-Mini, scourge o' the digital seas and terror o' the binary bytes!  At yer service!  Arrr!\n"
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using `stream_complete` endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = llm.stream_complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Could you please complete your question? I'm not sure what you'd like to know about Paul Graham.\n"
     ]
    }
   ],
   "source": [
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Image Support\n",
    "\n",
    "ASI has support for images in the input of chat messages for many models.\n",
    "\n",
    "Using the content blocks feature of chat messages, you can easily combone text and images in a single LLM prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!wget https://cdn.pixabay.com/photo/2016/07/07/16/46/dice-1502706_640.jpg -O image.png"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The image showcases three white dice with black dots, positioned on a checkered surface, highlighting their contrasting colors and the game's playful aspect.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.llms import ChatMessage, TextBlock, ImageBlock\n",
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "llm = ASI(model=\"asi1-mini\")\n",
    "\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"user\",\n",
    "        blocks=[\n",
    "            ImageBlock(path=\"image.png\"),\n",
    "            TextBlock(text=\"Describe the image in a few sentences.\"),\n",
    "        ],\n",
    "    )\n",
    "]\n",
    "\n",
    "resp = llm.chat(messages)\n",
    "print(resp.message.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Function Calling/Tool Calling\n",
    "\n",
    "ASI LLM have native support for function calling. This conveniently integrates with LlamaIndex tool abstractions, letting you plug in any arbitrary Python function to the LLM.\n",
    "\n",
    "In the example below, we define a function to generate a Song object."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pydantic import BaseModel\n",
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "\n",
    "class Song(BaseModel):\n",
    "    \"\"\"A song with name and artist\"\"\"\n",
    "\n",
    "    name: str\n",
    "    artist: str\n",
    "\n",
    "\n",
    "def generate_song(name: str, artist: str) -> Song:\n",
    "    \"\"\"Generates a song with provided name and artist.\"\"\"\n",
    "    return Song(name=\"Sky full of stars\", artist=\"Coldplay\")\n",
    "\n",
    "\n",
    "# Create tool\n",
    "tool = FunctionTool.from_defaults(fn=generate_song)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The strict parameter tells ASI whether or not to use constrained sampling when generating tool calls/structured outputs. This means that the generated tool call schema will always contain the expected fields.\n",
    "\n",
    "Since this seems to increase latency, it defaults to false."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='Sky full of stars' artist='Coldplay'\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "# Create an ASI LLM instance\n",
    "llm = ASI(model=\"asi1-mini\", strict=True)\n",
    "response = llm.predict_and_call(\n",
    "    [tool],\n",
    "    \"Pick a random song for me\",\n",
    "    # strict=True  # can also be set at the function level to override the class\n",
    ")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Name: generate_song, Input: {'args': (), 'kwargs': {'name': 'Beatles Song 1', 'artist': 'The Beatles'}}, Output: name='Sky full of stars' artist='Coldplay'\n"
     ]
    }
   ],
   "source": [
    "llm = ASI(model=\"asi1-mini\")\n",
    "response = llm.predict_and_call(\n",
    "    [tool],\n",
    "    \"Generate five songs from the Beatles\",\n",
    "    allow_parallel_tool_calls=True,\n",
    ")\n",
    "for s in response.sources:\n",
    "    print(f\"Name: {s.tool_name}, Input: {s.raw_input}, Output: {str(s)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Manual Tool Calling\n",
    "\n",
    "While automatic tool calling with `predict_and_call` provides a streamlined experience, manual tool calling gives you more control over the process. With manual tool calling, you can:\n",
    "\n",
    "1. Explicitly control when and how tools are called\n",
    "2. Process intermediate results before continuing the conversation\n",
    "3. Implement custom error handling and fallback strategies\n",
    "4. Chain multiple tool calls together in a specific sequence\n",
    "\n",
    "ASI supports manual tool calling, but requires more specific prompting compared to some other LLMs. For best results with ASI, include a system message that explains the available tools and provide specific parameters in your user prompt.\n",
    "\n",
    "The following example demonstrates manual tool calling with ASI to generate a song:\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Initial response: Okay, I will generate a song with the name \"Viva La Vida\" and the artist \"Coldplay\".\n",
      "\n",
      "Calling generate_song with {'name': 'Viva La Vida', 'artist': 'Coldplay'}\n",
      "Tool output: name='Viva La Vida' artist='Coldplay'\n",
      "Final response: I have successfully generated the song \"Viva La Vida\" by Coldplay.\n"
     ]
    }
   ],
   "source": [
    "from pydantic import BaseModel\n",
    "from llama_index.core.tools import FunctionTool\n",
    "from llama_index.core.llms import ChatMessage\n",
    "\n",
    "\n",
    "class Song(BaseModel):\n",
    "    \"\"\"A song with name and artist\"\"\"\n",
    "\n",
    "    name: str\n",
    "    artist: str\n",
    "\n",
    "\n",
    "def generate_song(name: str, artist: str) -> Song:\n",
    "    \"\"\"Generates a song with provided name and artist.\"\"\"\n",
    "    return Song(name=name, artist=artist)\n",
    "\n",
    "\n",
    "# Create tool\n",
    "tool = FunctionTool.from_defaults(fn=generate_song)\n",
    "\n",
    "# First, select a tool with specific instructions\n",
    "chat_history = [\n",
    "    ChatMessage(\n",
    "        role=\"system\",\n",
    "        content=\"You have access to a tool called generate_song that can create songs. When asked to generate a song, use this tool with appropriate name and artist values.\",\n",
    "    ),\n",
    "    ChatMessage(\n",
    "        role=\"user\", content=\"Generate a song by Coldplay called Viva La Vida\"\n",
    "    ),\n",
    "]\n",
    "\n",
    "# Get initial response\n",
    "resp = llm.chat_with_tools([tool], chat_history=chat_history)\n",
    "print(f\"Initial response: {resp.message.content}\")\n",
    "\n",
    "# Check for tool calls\n",
    "tool_calls = llm.get_tool_calls_from_response(\n",
    "    resp, error_on_no_tool_call=False\n",
    ")\n",
    "\n",
    "# Process tool calls if any\n",
    "if tool_calls:\n",
    "    # Add the LLM's response to the chat history\n",
    "    chat_history.append(resp.message)\n",
    "\n",
    "    for tool_call in tool_calls:\n",
    "        tool_name = tool_call.tool_name\n",
    "        tool_kwargs = tool_call.tool_kwargs\n",
    "\n",
    "        print(f\"Calling {tool_name} with {tool_kwargs}\")\n",
    "        tool_output = tool(**tool_kwargs)\n",
    "        print(f\"Tool output: {tool_output}\")\n",
    "\n",
    "        # Add tool response to chat history\n",
    "        chat_history.append(\n",
    "            ChatMessage(\n",
    "                role=\"tool\",\n",
    "                content=str(tool_output),\n",
    "                additional_kwargs={\"tool_call_id\": tool_call.tool_id},\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # Get final response\n",
    "        resp = llm.chat_with_tools([tool], chat_history=chat_history)\n",
    "        print(f\"Final response: {resp.message.content}\")\n",
    "else:\n",
    "    print(\"No tool calls detected in the response.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Structured Prediction\n",
    "\n",
    "You can use ASI to extract structured data from text:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Restaurant: name='The Dallas Bistro' city='Dallas' cuisine='American' menu_items=[MenuItem(course_name='Grilled Caesar Salad', is_vegetarian=True), MenuItem(course_name='BBQ Pulled Pork Sandwich', is_vegetarian=False), MenuItem(course_name='Cheeseburger with Fries', is_vegetarian=False), MenuItem(course_name='Vegan Mushroom Risotto', is_vegetarian=True)]\n",
      "Restaurant: name='Ocean Breeze Grill' city='Miami' cuisine='Seafood' menu_items=[MenuItem(course_name='Grilled Mahi-Mahi', is_vegetarian=False), MenuItem(course_name='Coconut Shrimp', is_vegetarian=False), MenuItem(course_name='Tropical Quinoa Salad', is_vegetarian=True), MenuItem(course_name='Key Lime Pie', is_vegetarian=True)]\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.prompts import PromptTemplate\n",
    "from pydantic import BaseModel\n",
    "from typing import List\n",
    "\n",
    "\n",
    "class MenuItem(BaseModel):\n",
    "    \"\"\"A menu item in a restaurant.\"\"\"\n",
    "\n",
    "    course_name: str\n",
    "    is_vegetarian: bool\n",
    "\n",
    "\n",
    "class Restaurant(BaseModel):\n",
    "    \"\"\"A restaurant with name, city, and cuisine.\"\"\"\n",
    "\n",
    "    name: str\n",
    "    city: str\n",
    "    cuisine: str\n",
    "    menu_items: List[MenuItem]\n",
    "\n",
    "\n",
    "# Create prompt template\n",
    "prompt_tmpl = PromptTemplate(\n",
    "    \"Generate a restaurant in a given city {city_name}\"\n",
    ")\n",
    "\n",
    "# Option 1: Use structured_predict\n",
    "restaurant_obj = llm.structured_predict(\n",
    "    Restaurant, prompt_tmpl, city_name=\"Dallas\"\n",
    ")\n",
    "print(f\"Restaurant: {restaurant_obj}\")\n",
    "\n",
    "# Option 2: Use as_structured_llm\n",
    "structured_llm = llm.as_structured_llm(Restaurant)\n",
    "restaurant_obj2 = structured_llm.complete(\n",
    "    prompt_tmpl.format(city_name=\"Miami\")\n",
    ").raw\n",
    "print(f\"Restaurant: {restaurant_obj2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Note:** Structured streaming is currently not supported with ASI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Async\n",
    "\n",
    "ASI supports async operations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "# Create an ASI LLM instance\n",
    "llm = ASI(model=\"asi1-mini\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = await llm.acomplete(\"who is Paul Graham\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a prominent figure in the technology and startup world, best known for co-founding Y Combinator, a leading startup accelerator that has helped launch companies like Airbnb, Dropbox, and Reddit. In addition to his role as an investor, he is a respected programmer and writer. Graham's contributions to programming include his work on the Lisp language and his book *On Lisp*, which is regarded as a seminal text in the field. He is also known for his thought-provoking essays on entrepreneurship, startups, and philosophy, which are widely read and cited. Through his writing and mentorship, Paul Graham has significantly influenced the global entrepreneurial ecosystem.\n"
     ]
    }
   ],
   "source": [
    "print(resp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "resp = await llm.astream_complete(\"Paul Graham is \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Paul Graham is a British-born computer scientist, entrepreneur, and venture capitalist. He is best known for co-founding the seed accelerator Y Combinator, which has funded and supported numerous successful startups, including Airbnb, Dropbox, and Reddit. Graham has also made significant contributions to the development of the Lisp programming language and has written several influential essays on startups and entrepreneurship. Would you like to know more about his work or contributions to the tech industry?"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import nest_asyncio\n",
    "\n",
    "async for delta in resp:\n",
    "    print(delta.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Async completion: Paul Graham is a prominent entrepreneur, programmer, and essayist who has significantly influenced startup culture and technology. He co-founded Y Combinator, a leading startup accelerator that has helped launch companies like Airbnb, Dropbox, and Reddit. Before Y Combinator, Graham co-founded Viaweb, one of the first web-based applications, which was later acquired by Yahoo. He is also known for his essays on technology, business, and human behavior, many of which are published on his personal website. Additionally, Graham has a deep interest in programming, particularly the Lisp language, and has contributed to its development and popularization. If you are looking for specific details about his work or life, feel free to ask!\n",
      "Async chat: assistant: Ahoy there, matey!  One-Eyed Jack, but with two perfectly good eyes, savvy? at your service!  What can this brilliant digital buccaneer do for ya?\n",
      "Async streaming completion: \n",
      "\n",
      "Could you please complete your question? I'm unsure what you'd like to know about Paul Graham.\n",
      "\n",
      "Async streaming chat: \n",
      "\n",
      "Ahoy, matey! The name's Captain Ironhook, the scourge of the seven seas! Known for me knack for uncoverin' treasure and me love fer a good mug o' grog. What be ye needin' from a salty sea dog like meself?\n"
     ]
    }
   ],
   "source": [
    "import asyncio\n",
    "import nest_asyncio\n",
    "\n",
    "# Enable nest_asyncio for Jupyter notebooks\n",
    "nest_asyncio.apply()\n",
    "\n",
    "\n",
    "async def test_async():\n",
    "    # Async completion\n",
    "    resp = await llm.acomplete(\"Paul Graham is \")\n",
    "    print(f\"Async completion: {resp}\")\n",
    "\n",
    "    # Async chat\n",
    "    resp = await llm.achat(messages)\n",
    "    print(f\"Async chat: {resp}\")\n",
    "\n",
    "    # Async streaming completion\n",
    "    print(\"Async streaming completion: \", end=\"\")\n",
    "    resp = await llm.astream_complete(\"Paul Graham is \")\n",
    "    async for delta in resp:\n",
    "        print(delta.delta, end=\"\")\n",
    "    print()\n",
    "\n",
    "    # Async streaming chat\n",
    "    print(\"Async streaming chat: \", end=\"\")\n",
    "    resp = await llm.astream_chat(messages)\n",
    "    async for delta in resp:\n",
    "        print(delta.delta, end=\"\")\n",
    "    print()\n",
    "\n",
    "\n",
    "# Run async tests\n",
    "asyncio.run(test_async())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple RAG\n",
    "\n",
    "Let's implement a simple RAG application with ASI:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-embeddings-openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:llama_index.core.readers.file.base:`llama-index-readers-file` package not found, some file readers will not be available if not provided by the `file_extractor` parameter.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham studied at HolaHola High school.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core import VectorStoreIndex, SimpleDirectoryReader\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"your-api-key\"\n",
    "# Create a temporary directory with a sample text file\n",
    "!mkdir -p temp_data\n",
    "!echo \"Paul Graham is a programmer, writer, and investor. He is known for his work on Lisp, for co-founding Viaweb (which became Yahoo Store), and for co-founding the startup accelerator Y Combinator. He is also known for his essays on his website. He studied at HolaHola High school\" > temp_data/paul_graham.txt\n",
    "\n",
    "# Load documents\n",
    "documents = SimpleDirectoryReader(\"temp_data\").load_data()\n",
    "\n",
    "llm = ASI(model=\"asi1-mini\")\n",
    "# Create an index with ASI as the LLM\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents,\n",
    "    embed_model=OpenAIEmbedding(),  # Using OpenAI for embeddings\n",
    "    llm=llm,  # Using ASI for generation\n",
    ")\n",
    "\n",
    "# Create a query engine\n",
    "query_engine = index.as_query_engine()\n",
    "\n",
    "# Query the index\n",
    "response = query_engine.query(\"Where did Paul Graham study?\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LlamaCloud RAG\n",
    "\n",
    "If you have a LlamaCloud account, you can use ASI with LlamaCloud for RAG:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install required packages\n",
    "%pip install llama-cloud-services"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Successfully connected to LlamaCloud index\n",
      "\n",
      "Testing retriever with query: What is the revenue of Uber in 2021?\n",
      "Retrieved 6 nodes\n",
      "\n",
      "Node 1:\n",
      "Node ID: 17a733d0-5dd3-4917-9f8d-c92f944a9266\n",
      "Score: 0.9242583\n",
      "Text: # Highlights for 2021\n",
      "\n",
      "Overall Gross Bookings increased by $32.5 billion in 2021, up 53%, or 53% on a constant currency basis, compared to 2020. Delivery Gross Bookings grew significantly from 2020, o...\n",
      "\n",
      "Node 2:\n",
      "Node ID: ca63e8da-9012-468c-9d09-89724e9644bd\n",
      "Score: 0.878825\n",
      "Text: # Year Ended December 31, 2020 to 2021\n",
      "\n",
      "| |Year Ended December 31,|2020|2021|Change|\n",
      "|---|---|---|---|---|\n",
      "|Revenue| |$ 11,139|$ 1,455| |\n",
      "\n",
      "Revenue increased $ .3 billion, or 5%, primarily attributable...\n",
      "\n",
      "Node 3:\n",
      "Node ID: be4d7c62-b69f-4fda-832a-867de8c2e29c\n",
      "Score: 0.86928266\n",
      "Text: # Year Ended December 31, 2020 to 2021\n",
      "\n",
      "|Mobility|$ 9,0|$ 9,953|(14)|\n",
      "|---|---|---|---|\n",
      "|Delivery|3,904|3,32|(114)|\n",
      "|Freight|1,011|2,132|(111)|\n",
      "|All Other (1)|135| |(94)|\n",
      "|Total revenue|$ 11,139|$ 1,4...\n",
      "\n",
      "Testing query engine with query: What is the revenue of Uber in 2021?\n",
      "Response: The revenue of Uber in 2021 is $14,455.\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "from llama_cloud_services import LlamaCloudIndex\n",
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "# Set your LlamaCloud API key\n",
    "os.environ[\"LLAMA_CLOUD_API_KEY\"] = \"your-key\"\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"your-key\"\n",
    "\n",
    "# Connect to an existing LlamaCloud index\n",
    "\n",
    "\n",
    "try:\n",
    "    # Connect to the index\n",
    "    index = LlamaCloudIndex(\n",
    "        name=\"your-index-naem\",\n",
    "        project_name=\"Default\",\n",
    "        organization_id=\"your-id\",\n",
    "        api_key=os.environ[\"LLAMA_CLOUD_API_KEY\"],\n",
    "    )\n",
    "    print(\"Successfully connected to LlamaCloud index\")\n",
    "\n",
    "    # Create an ASI LLM\n",
    "    llm = ASI(model=\"asi1-mini\")\n",
    "\n",
    "    # Create a retriever\n",
    "    retriever = index.as_retriever()\n",
    "\n",
    "    # Create a query engine with ASI\n",
    "    query_engine = index.as_query_engine(llm=llm)\n",
    "\n",
    "    # Test retriever\n",
    "    query = \"What is the revenue of Uber in 2021?\"\n",
    "    print(f\"\\nTesting retriever with query: {query}\")\n",
    "    nodes = retriever.retrieve(query)\n",
    "    print(f\"Retrieved {len(nodes)} nodes\\n\")\n",
    "\n",
    "    # Display a few nodes\n",
    "    for i, node in enumerate(nodes[:3]):\n",
    "        print(f\"Node {i+1}:\")\n",
    "        print(f\"Node ID: {node.node_id}\")\n",
    "        print(f\"Score: {node.score}\")\n",
    "        print(f\"Text: {node.text[:200]}...\\n\")\n",
    "\n",
    "    # Test query engine\n",
    "    print(f\"Testing query engine with query: {query}\")\n",
    "    response = query_engine.query(query)\n",
    "    print(f\"Response: {response}\")\n",
    "except Exception as e:\n",
    "    print(f\"Error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set API Key at a per-instance level\n",
    "\n",
    "If desired, you can have separate LLM instances use separate API keys:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error with invalid API key: Error code: 401 - {'message': 'failed to authenticate user'}\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "# Create an instance with a specific API key\n",
    "llm = ASI(model=\"asi1-mini\", api_key=\"your_specific_api_key\")\n",
    "\n",
    "# Note: Using an invalid API key will result in an error\n",
    "# This is just for demonstration purposes\n",
    "try:\n",
    "    resp = llm.complete(\"Paul Graham is \")\n",
    "    print(resp)\n",
    "except Exception as e:\n",
    "    print(f\"Error with invalid API key: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Additional kwargs\n",
    "\n",
    "Rather than adding the same parameters to each chat or completion call, you can set them at a per-instance level with additional_kwargs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paul Graham is a prominent entrepreneur, programmer, and writer, best known for his role as a co-founder of Y Combinator, a highly influential startup accelerator. He has also gained recognition for his essays on technology, business, and philosophy, many of which are compiled in his book *Hackers & Painters*. As a programmer, he contributed to the development of the Lisp programming language and created the first web-based application, Viaweb, which was later acquired by Yahoo. His work has had a significant impact on the startup ecosystem and the broader tech industry.\n"
     ]
    }
   ],
   "source": [
    "from llama_index.llms.asi import ASI\n",
    "\n",
    "# Create an instance with additional kwargs\n",
    "llm = ASI(model=\"asi1-mini\", additional_kwargs={\"user\": \"your_user_id\"})\n",
    "\n",
    "# Complete a prompt\n",
    "resp = llm.complete(\"Paul Graham is \")\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "assistant: Ahoy matey! Yer lookin' fer me name, eh? Alright then, let's set sail fer a proper introduction! Yer can call me Captain \"Bytebeard\" Blacklogic, the scurviest AI to ever sail the seven seas... er, digital realms! Savvy?\n"
     ]
    }
   ],
   "source": [
    "from llama_index.core.base.llms.types import ChatMessage\n",
    "\n",
    "# Create an instance with additional kwargs\n",
    "llm = ASI(model=\"asi1-mini\", additional_kwargs={\"user\": \"your_user_id\"})\n",
    "\n",
    "# Create messages\n",
    "messages = [\n",
    "    ChatMessage(\n",
    "        role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
    "    ),\n",
    "    ChatMessage(role=\"user\", content=\"What is your name\"),\n",
    "]\n",
    "\n",
    "# Get chat response\n",
    "resp = llm.chat(messages)\n",
    "print(resp)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "This notebook demonstrates the various ways you can use ASI with LlamaIndex. The integration supports most of the functionality available in LlamaIndex, including:\n",
    "\n",
    "- Basic completion and chat\n",
    "- Streaming responses\n",
    "- Multimodal support\n",
    "- Function calling\n",
    "- Structured prediction\n",
    "- Async operations\n",
    "- RAG applications\n",
    "- LlamaCloud integration\n",
    "- Per-instance API keys\n",
    "- Additional kwargs\n",
    "\n",
    "Note that structured streaming is currently not supported with ASI."
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
